View Javadoc
1   package edu.jiangxin.apktoolbox.convert.protobuf.supervised;
2   
3   import com.google.protobuf.DescriptorProtos;
4   import com.google.protobuf.Descriptors;
5   import com.google.protobuf.InvalidProtocolBufferException;
6   
7   import java.io.IOException;
8   import java.io.UncheckedIOException;
9   import java.nio.file.Files;
10  import java.nio.file.Path;
11  import java.util.*;
12  import java.util.stream.Stream;
13  
14  public final class DescriptorCache {
15      private static final Descriptors.FileDescriptor[] DEPENDENCIES = new Descriptors.FileDescriptor[0];
16  
17      public static DescriptorCache emptyCache() {
18          return new DescriptorCache();
19      }
20  
21      public static DescriptorCache fromDirectory(final Path directory) {
22          Objects.requireNonNull(directory);
23  
24          if (!Files.isDirectory(directory)) {
25              throw new IllegalArgumentException("Path must be a directory: " + directory);
26          }
27  
28          final DescriptorCache cache = new DescriptorCache();
29          try (Stream<Path> walk = Files.walk(directory)) {
30              walk.filter(Files::isRegularFile)
31                      .forEach(cache::addDescriptors);
32          } catch (final IOException e) {
33              throw new UncheckedIOException(e);
34          }
35          return cache;
36      }
37  
38      public static DescriptorCache fromFile(final Path descriptorsFile) {
39          Objects.requireNonNull(descriptorsFile);
40  
41          if (!Files.isRegularFile(descriptorsFile)) {
42              throw new IllegalArgumentException("Path must be a regular file: " + descriptorsFile);
43          }
44          final DescriptorCache cache = new DescriptorCache();
45          cache.addDescriptors(descriptorsFile);
46          return cache;
47      }
48  
49      private final Map<String, Descriptors.Descriptor> typeNameToDescriptor = new HashMap<>();
50  
51      private DescriptorCache() {
52  
53      }
54  
55      @SuppressWarnings({"WeakerAccess", "UnusedReturnValue"})
56      public Optional<Descriptors.Descriptor> addDescriptor(final Descriptors.Descriptor descriptor) {
57          Objects.requireNonNull(descriptor);
58          final String typeName = Objects.requireNonNull(descriptor.getName());
59  
60          return Optional.ofNullable(typeNameToDescriptor.put(typeName, descriptor));
61      }
62  
63      @SuppressWarnings("WeakerAccess")
64      public void addDescriptors(final Path descriptorsFile) {
65          Objects.requireNonNull(descriptorsFile);
66  
67          if (!Files.isRegularFile(descriptorsFile)) {
68              throw new IllegalArgumentException("Path must be a regular file: " + descriptorsFile);
69          }
70  
71          try {
72              addDescriptors(Files.readAllBytes(descriptorsFile));
73          } catch (final IOException e) {
74              throw new UncheckedIOException("While reading: " + descriptorsFile.toAbsolutePath(), e);
75          }
76      }
77  
78      @SuppressWarnings("WeakerAccess")
79      public void addDescriptors(final byte[] descriptorsRaw) {
80          Objects.requireNonNull(descriptorsRaw);
81  
82          try {
83              final DescriptorProtos.FileDescriptorSet descriptorSet =
84                      DescriptorProtos.FileDescriptorSet.parseFrom(descriptorsRaw);
85              for (final DescriptorProtos.FileDescriptorProto descriptorFile : descriptorSet.getFileList()) {
86                  final Descriptors.FileDescriptor fileDescriptor =
87                          Descriptors.FileDescriptor.buildFrom(descriptorFile, DEPENDENCIES);
88                  for (final Descriptors.Descriptor descriptor : fileDescriptor.getMessageTypes()) {
89                      addDescriptor(descriptor);
90                  }
91              }
92          } catch (final InvalidProtocolBufferException | Descriptors.DescriptorValidationException e) {
93              throw new RuntimeException(e);
94          }
95      }
96  
97      public Optional<Descriptors.Descriptor> getByTypeName(final String typeName) {
98          Objects.requireNonNull(typeName);
99  
100         return Optional.ofNullable(typeNameToDescriptor.get(typeName));
101     }
102 
103     public Collection<Descriptors.Descriptor> getDescriptors() {
104         return Collections.unmodifiableCollection(typeNameToDescriptor.values());
105     }
106 
107     public Collection<Map.Entry<String, Descriptors.Descriptor>> getEntries() {
108         return Collections.unmodifiableCollection(typeNameToDescriptor.entrySet());
109     }
110 
111     public boolean isEmpty() {
112         return typeNameToDescriptor.isEmpty();
113     }
114 
115     public int size() {
116         return typeNameToDescriptor.size();
117     }
118 }